#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# 本地模型部署工具

import os
import sys
import argparse
import logging
from pathlib import Path
from funasr.utils.local_model_manager import get_local_model_manager

def setup_logging(level=logging.INFO):
    """设置日志级别"""
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(sys.stdout)
        ]
    )

def list_models():
    """列出所有本地模型"""
    manager = get_local_model_manager()
    models = manager.list_local_models()
    
    if not models:
        print("没有本地部署的模型")
        return
        
    print("\n本地模型列表:")
    print("-" * 60)
    for model_name, model_type in models.items():
        info = manager.get_model_info(model_name)
        if info:
            size_mb = info.get('size', 0) / (1024 * 1024)
            print(f"模型名称: {model_name}")
            print(f"模型类型: {model_type}")
            print(f"本地路径: {info.get('path', 'N/A')}")
            print(f"模型大小: {size_mb:.2f} MB")
            print(f"文件数量: {len(info.get('files', []))}")
            print("-" * 60)
        else:
            print(f"模型: {model_name} ({model_type})")

def deploy_model(model_path, model_name):
    """部署模型到本地"""
    if not os.path.exists(model_path):
        print(f"错误: 模型路径不存在: {model_path}")
        return False
        
    manager = get_local_model_manager()
    success = manager.setup_model(model_name, model_path)
    
    if success:
        print(f"✓ 模型 {model_name} 部署成功!")
        print(f"  本地路径: {manager.get_model_path(model_name)}")
    else:
        print(f"✗ 模型 {model_name} 部署失败!")
    
    return success

def check_model(model_name):
    """检查模型是否在本地可用"""
    manager = get_local_model_manager()
    available = manager.is_model_available(model_name)
    
    if available:
        info = manager.get_model_info(model_name)
        print(f"✓ 模型 {model_name} 已在本地部署")
        if info:
            size_mb = info.get('size', 0) / (1024 * 1024)
            print(f"  本地路径: {info.get('path', 'N/A')}")
            print(f"  模型大小: {size_mb:.2f} MB")
    else:
        print(f"✗ 模型 {model_name} 未在本地部署")
        print("  将使用远程模型（需要下载）")
    
    return available

def remove_model(model_name):
    """删除本地模型"""
    manager = get_local_model_manager()
    
    if not manager.is_model_available(model_name):
        print(f"模型 {model_name} 不存在")
        return False
        
    confirm = input(f"确定要删除模型 {model_name} 吗? (y/N): ")
    if confirm.lower() != 'y':
        print("取消删除操作")
        return False
        
    success = manager.remove_model(model_name)
    
    if success:
        print(f"✓ 模型 {model_name} 已删除")
    else:
        print(f"✗ 删除模型 {model_name} 失败")
    
    return success

def main():
    parser = argparse.ArgumentParser(description='FunASR 本地模型管理工具')
    parser.add_argument('--verbose', '-v', action='store_true', help='显示详细日志')
    
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # list 命令
    list_parser = subparsers.add_parser('list', help='列出本地模型')
    
    # deploy 命令
    deploy_parser = subparsers.add_parser('deploy', help='部署模型到本地')
    deploy_parser.add_argument('model_path', help='模型源目录路径')
    deploy_parser.add_argument('model_name', help='模型名称或别名')
    
    # check 命令
    check_parser = subparsers.add_parser('check', help='检查模型是否可用')
    check_parser.add_argument('model_name', help='模型名称或别名')
    
    # remove 命令
    remove_parser = subparsers.add_parser('remove', help='删除本地模型')
    remove_parser.add_argument('model_name', help='模型名称或别名')
    
    args = parser.parse_args()
    
    if args.verbose:
        setup_logging(logging.DEBUG)
    else:
        setup_logging(logging.INFO)
    
    if not args.command:
        parser.print_help()
        return
    
    if args.command == 'list':
        list_models()
    elif args.command == 'deploy':
        deploy_model(args.model_path, args.model_name)
    elif args.command == 'check':
        check_model(args.model_name)
    elif args.command == 'remove':
        remove_model(args.model_name)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()